import torch
import torch.nn as nn
import torch.nn.functional as F

class Lossor(object):
    def __init__(self, focal_weight):
        self.sum_loss = 0.0
        self.num = 0
        self.focal_weight = focal_weight

    def MixLoss(self, out, label, add=True):
        b1, b2, b3, b4, fuse, r1 = out
        label = label.long()
        s1 = F.cross_entropy(b1, label, ignore_index=255)
        s2 = F.cross_entropy(b2, label, ignore_index=255)
        s3 = F.cross_entropy(b3, label, ignore_index=255)
        s4 = F.cross_entropy(b4, label, ignore_index=255)
        sf = F.cross_entropy(fuse, label, ignore_index=255)
        sl = s1 + s2 + s3 + s4 + sf
        fl = self.focal_loss(r1, label)
        total_loss = sl + fl * self.focal_weight
        if add:
            self.sum_loss += total_loss.item()
            self.num += 1
        return total_loss

    def mean_loss(self, clear=True):
        loss = self.sum_loss / self.num
        if clear:
            self.sum_loss = 0.0
            self.num_loss = 0
        return loss

    def focal_loss(self, pred, label, alpha=0.25, gamma=2.0):
        label = self.flat_dim(label, pred.size()[1]).float()
        pk = torch.sum(pred * label, 1)
        loss = - alpha * torch.mean(torch.pow(1.0 - pk, gamma) * torch.log(torch.clamp(pk, 1e-12, 1.0)))
        return loss

    def flat_dim(self, label, num_class):
        out = []
        for i in range(num_class):
            out.append(label == i)
        out = torch.stack(out, dim=1)
        return out


class Evaluator(object):
    def __init__(self, num_class):
        self.num_class = num_class
        self.matrix = torch.zeros((self.num_class, self.num_class))

    def add_batch(self, gt_image, pre_image):
        assert gt_image.shape == pre_image.shape
        self.matrix += self.generate_matrix(gt_image, pre_image).cpu()

    def mIoU(self, clear=False):
        iou = torch.diag(self.matrix) / (
                torch.sum(self.matrix, 0) + torch.sum(self.matrix, 1) - torch.diag(self.matrix))

        miou = torch.mean(iou[torch.isfinite(iou)])
        if clear:
            self.matrix = torch.zeros((self.num_class, self.num_class))
        return miou.item()

    def acc(self, clear=False):
        acc = torch.sum(torch.diag(self.matrix))/torch.sum(self.matrix)
        if clear:
            self.matrix = torch.zeros((self.num_class, self.num_class))
        return acc

    def generate_matrix(self, gt_image, pre_image):
        gt_image = gt_image.data.float()
        pre_image = pre_image.data.float()
        mask = (gt_image >= 0) & (gt_image < self.num_class)
        label = self.num_class * gt_image[mask] + pre_image[mask]
        count = torch.bincount(label.int(), minlength=self.num_class ** 2)
        confusion_matrix = count[:self.num_class ** 2].reshape(self.num_class, self.num_class)  # int can't div 0
        return confusion_matrix.float()